//
//  SSHTunnel.m
//  tunneltool
//
//  Created by Andreas Hindborg on 17/05/11.
//  Copyright 2011 Andreas Hindborg. All rights reserved.
//

#import "SSHTunnel.h"
#include <stdio.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/select.h>
#include <libssh2.h>
#include <string.h>
#import "SSHSession.h"

@interface SSHTunnel ()
@property (readwrite, retain) NSString *remoteIP;
@property (readwrite, retain) NSString *localIP;
@property (readwrite, assign)  uint32_t remotePort;
@property (readwrite, assign)  uint16_t localPort;
@property (readwrite, retain)  SSHSession *session;
@property (readwrite, assign)  bool connected;
@property (readwrite, retain)  NSMutableArray *connections;
@end

@implementation SSHTunnel

@synthesize remoteIP;
@synthesize localIP;
@synthesize remotePort;
@synthesize localPort;
@synthesize connected;
@synthesize session;
@synthesize connections;

- (SSHTunnel*) initWithSession: (SSHSession*) aSession
                      remoteIP: (NSString*)   aRemoteIP
                    remotePort: (uint32_t)    aRemotePort
                       localIP: (NSString*)   aLocalIP
                     localPort: (uint32_t)    aLocalPort
{
    self = [super init];
    if(!self) return self;

    [self setRemoteIP: aRemoteIP];
    [self setLocalIP: aLocalIP];
    [self setRemotePort: aRemotePort];
    [self setLocalPort: aLocalPort];
    [self setSession: aSession];
    [self setConnected: false];
    [self setConnections: [[NSMutableArray alloc] init]];
    shouldDie = false;
    
    return self;
}

- (int) activate
{
    struct sockaddr_in listenAddress;
    int i;
    int err;
    
	memset(&listenAddress, 0, sizeof(listenAddress));

    if(![session connected]) {
        NSLog(@"SSHTunnel: activate: Failed. Session is not connected!");
        return 10;
    }
	
    NSLog(@"Activating connection");

    // Create local listen socket
    localListenSocket = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
    
    // Setup address
    listenAddress.sin_family = AF_INET;
    listenAddress.sin_port = htons(localPort);
    listenAddress.sin_addr.s_addr = inet_addr([localIP UTF8String]);

    if(listenAddress.sin_addr.s_addr == INADDR_NONE) {
        NSLog(@"SSHTunnel: Invalid local IP: %@!", localIP);
        return 1;
    }

    // Set reuse address option
    i = 1;
    err = setsockopt(localListenSocket, SOL_SOCKET, SO_REUSEADDR, &i, sizeof(i));
    if(err) {
        NSLog(@"setsockopt: %s", strerror(errno));
        return 2;
    }
    
    // Bind to local address
    err = bind(localListenSocket, (struct sockaddr*) &listenAddress, sizeof(listenAddress));
    if(err == -1) {
        NSLog(@"bind: %s", strerror(errno));
        NSLog(@"localIP: %s", inet_ntoa(listenAddress.sin_addr));
        NSLog(@"localPort: %d", ntohs(listenAddress.sin_port));
        NSLog(@"errorCode: %d", errno);
        return 3;
    }
    
    // Listen for incomming connections
    err = listen(localListenSocket, 1);
    if(err) {
        NSLog(@"listen: %s", strerror(errno));
        return 4;
    }
    
    connected = true;

    // Dispatch accept loop
    [self performSelectorInBackground:@selector(serve) withObject:nil];
    
    NSLog(@"Connection activated");
    return 0;

}

- (void) serve
{
    NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init];

    LIBSSH2_SESSION *libsshSession = [session getSession];

    // Loop untill we are killed
    while(!shouldDie) {
        fd_set readSet;
        int selectedDescriptors;
        struct timeval timeout ={0, 100000};
        NSEnumerator *connEnum;
        id connection;
        int maxDescriptor = localListenSocket;

        FD_ZERO(&readSet);
        FD_SET(localListenSocket, &readSet);
        FD_SET([session sessionSocket], &readSet);

        //TODO: Put other listen sockets in read set
        connEnum = [connections objectEnumerator];
        while(connection = [connEnum nextObject]) {
            int connectionDescriptor = [connection clientSocket];
            FD_SET(connectionDescriptor, &readSet);
            maxDescriptor = connectionDescriptor > maxDescriptor ? connectionDescriptor : maxDescriptor;
        }
        
        // Select
        selectedDescriptors = select(maxDescriptor+1, &readSet, NULL, NULL, &timeout);
        if(selectedDescriptors == -1) {
            NSLog(@"select: %s", strerror(errno));
            break;
        }

        
        libssh2_session_set_blocking(libsshSession, 0);

        // Iterate connections
        int length;
        int written;
        int partWritten;
        char buffer[2 << 13];

        connEnum = [connections objectEnumerator];
        while(connection = [connEnum nextObject]) {
            int connectionDescriptor = [connection clientSocket];
            LIBSSH2_CHANNEL *connectionChannel = [connection channel];
            
            // Read local sockets
            if(selectedDescriptors && FD_ISSET(connectionDescriptor, &readSet)) {
                //NSLog(@"Reading from descriptor %d", connectionDescriptor);

                length = recv(connectionDescriptor, buffer, sizeof(buffer), 0);
                if(length < 0) {
                    NSLog(@"recv: %s", strerror(errno));
                    break;
                }
                else if(length == 0) {
                    NSLog(@"Client disconnected: %@:%u", [connection clientHost], [connection clientPort]);
                    [connection setDead:true];
                    continue;
                }
                
                written = 0;
                while (written < length) {
                    partWritten = libssh2_channel_write(connectionChannel, buffer + written, length - written);
                    if (partWritten == LIBSSH2_ERROR_EAGAIN) break;
                    else if (partWritten < 0) {
                        NSLog(@"libssh2_channel_read: %d!", partWritten);
                        break; //TODO break out of 2 loops
                    }

                    written += partWritten;
                }
                //NSLog(@"Transferred %d bytes to channel", written);
            }

            // Read ssh channels
            length = libssh2_channel_read(connectionChannel, buffer, sizeof(buffer));
            if(length == LIBSSH2_ERROR_EAGAIN) continue; //TODO skip this channel
            else if (length == LIBSSH2_ERROR_CHANNEL_CLOSED) {
                NSLog(@"Client disconnected: %@:%u", [connection clientHost], [connection clientPort]);
                [connection setDead:true];
                continue;
            }
            else if (length < 0) {
                NSLog(@"libssh2_channel_read: %d!", length);
                return;
            }
            else if(length > 0) {
                //NSLog(@"Reading from ssh channel for descriptor %d", [connection clientSocket]);
                written = 0;
                while (written < length) {
                    partWritten = send(connectionDescriptor, buffer + written, length - written, 0);
                    if (partWritten <= 0) {
                        NSLog(@"send: %s", strerror(errno));
                        return;
                    }

                    written += partWritten;
                }


                //NSLog(@"Transferred %d bytes from channel", written);
            }
        }

        // Read from listen socket?
        if(selectedDescriptors && FD_ISSET(localListenSocket, &readSet)) {
            struct sockaddr_in clientAddress;
            socklen_t addressLength = sizeof(clientAddress);
            int forwardSocket;
            const char *clientHost;
            uint32_t clientPort;
            
            libssh2_session_set_blocking(libsshSession, 1);

            // Accept
            forwardSocket = accept(localListenSocket, (struct sockaddr*) &clientAddress, &addressLength);
            if(forwardSocket == -1) {
                NSLog(@"accept: %s", strerror(errno));
                break;
            }
            
            clientHost = inet_ntoa(clientAddress.sin_addr);
            clientPort = ntohs(clientAddress.sin_port);

            NSLog(@"Accepted connection from %s:%d (descriptor: %d)", clientHost, clientPort, forwardSocket);

            // Create tcp/ip channel
            LIBSSH2_CHANNEL *channel;
            const char *remoteIP_UTF8 = [remoteIP UTF8String];
            channel = libssh2_channel_direct_tcpip_ex(libsshSession, 
													  remoteIP_UTF8, 
													  remotePort, 
													  clientHost, 
													  clientPort);

            if(!channel) {
                NSLog(@"libssh2_channel_direct_tcpip_ex: Failed!");
                break;
            }

            SSHTunnelConnection * conn = [[SSHTunnelConnection alloc] init];
            conn.clientSocket = forwardSocket;
            conn.channel = channel;
            conn.clientHost = [NSString stringWithUTF8String:clientHost];
            conn.clientPort = clientPort;

            [connections addObject:conn];
            [conn release];

            NSLog(@"Channel connected on remtoe side to: %s:%u", remoteIP_UTF8, remotePort);

        }
        
        for (int i = 0; i < connections.count; i++) {
            SSHTunnelConnection *connection = [connections objectAtIndex:i];
            if (connection.dead) {
                i--;
                [connections removeObject:connection];
            }
        }
    }
    
    [pool drain];
    return;
}

@end
